package leetcode;

import java.util.*;

/**
 * @Project : algorithm-learning-java
 * @Package : leetcode
 * @Author : Ruoyu Wang
 * @User : Momenta
 * @DateTime : 2024/1/4 16:53
 */
public class Q0373 {

    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        PriorityQueue<ArrayList<Integer>> pq = new PriorityQueue<>((o1, o2) -> nums1[o1.get(0)] + nums2[o1.get(1)] - nums1[o2.get(0)] - nums2[o2.get(1)]);
        HashSet<String> set = new HashSet<>();
        List<List<Integer>> res = new ArrayList<>();

        pq.add(new ArrayList<>(Arrays.asList(0, 0)));
        set.add("0,0");
        while (res.size() < k) {
            ArrayList<Integer> tmp = pq.poll();
            set.remove(tmp.get(0) + "," + tmp.get(1));
            int i = tmp.get(0) + 1;
            int j = tmp.get(1) + 1;

            while (i < nums1.length && nums1[i] == nums1[i - 1]) {
                i++;
            }
            while (j < nums2.length && nums2[j] == nums2[j - 1]) {
                j++;
            }

            for (long l = 0; l < (long) (i - tmp.get(0)) * (j - tmp.get(1)); l++) {
                if (res.size() >= k) break;
                res.add(new ArrayList<>(Arrays.asList(nums1[tmp.get(0)], nums2[tmp.get(1)])));
            }

            if (i < nums1.length && !set.contains(i + "," + tmp.get(1))) {
                pq.add(new ArrayList<>(Arrays.asList(i, tmp.get(1))));
                set.add(i + "," + tmp.get(1));
            }

            if (j < nums2.length && !set.contains(tmp.get(0) + "," + j)) {
                pq.add(new ArrayList<>(Arrays.asList(tmp.get(0), j)));
                set.add(tmp.get(0) + "," + j);
            }
        }
        return res;
    }

    public static void main(String[] args) {
        Q0373 sl = new Q0373();
        int[] nums1 = {1, 2, 4, 5, 6};
        int[] nums2 = {3, 5, 7, 9};
        int k = 20;
        System.out.println(sl.kSmallestPairs(nums1, nums2, k));
    }
}
